/*
Copyright (C) 2022-2024 Inspur Corp.

SPDX-License-Identifier: Apache-2.0
*/
package service

import (
	"context"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/chromedp/cdproto/emulation"
	"github.com/chromedp/cdproto/network"
	"github.com/chromedp/cdproto/page"
	"github.com/chromedp/cdproto/runtime"
	"github.com/chromedp/chromedp"
	"github.com/go-xorm/xorm"
	"github.com/panjf2000/ants/v2"
	"github.com/pkg/errors"

	"icep/bean/param"
	"icep/chain/client"
	"icep/common/constant"
	"icep/common/crypto/gm"
	"icep/common/db"
	"icep/common/upload"
	"icep/common/utils"
	"icep/config"
)

var (
	//任务暂存队列
	taskQueue *TaskQueue
	queueLock sync.RWMutex

	//chrome context
	chromeCtx       context.Context
	chromeCtxCancel context.CancelFunc

	//任务管理器
	taskMana *TaskManager
	manaLock sync.RWMutex

	//格式化滚动页面的语句
	fmtScrollExpr = func(val int) string { return fmt.Sprintf("window.scrollTo(0,%s)", strconv.Itoa(val)) }
)

//取证任务的基本单元
type task struct {
	userId      string                           //用户ID
	evidenceId  string                           //取证ID
	identityId  string                           //身份ID
	userPrvKey  string                           //用户私钥
	para        *param.CreateObtainEvidenceParam //取证的请求参数
	taskType    string                           //取证的类型
	client      *client.IchainClient
	letterParam *param.EvidenceLetterParam //存证函参数
}

// TaskQueue 暂存网页取证任务的队列
type TaskQueue struct {
	length      int
	taskCacheCh chan *task
}

//初始化
func init() {
	//初始化任务队列
	GetTaskQueue()
	//初始化任务管理器
	mana, err := InitTaskManager()
	if err != nil {
		log.Errorf("init task manager error: %v", err)
		return
	}
	//启动任务管理器
	mana.startTask()

	//初始化chrome
	_, _, err = GetChromeCtx()
	if err != nil {
		log.Errorf("get chrome ctx error: %v", err)
		return
	}
}

// GetTaskQueue 初始化队列
func GetTaskQueue() *TaskQueue {
	//队列已初始化
	if taskQueue != nil {
		return taskQueue
	}
	//队列未初始化
	taskQueue = &TaskQueue{ //容量为取证并发数的200倍
		taskCacheCh: make(chan *task, config.GetInt("icep.cdp.task-capacity", 50)*200),
	}
	return taskQueue
}

// Push 将取证任务添加到缓存通道
func (queue *TaskQueue) Push(userId, evidenceId, userPrvKey string, para *param.CreateObtainEvidenceParam,
	taskType string, client *client.IchainClient, letterParam *param.EvidenceLetterParam) {
	queue.taskCacheCh <- &task{
		userId:      userId,
		evidenceId:  evidenceId,
		userPrvKey:  userPrvKey,
		para:        para,
		taskType:    taskType,
		client:      client,
		letterParam: letterParam,
	}
	queueLock.Lock()
	defer queueLock.Unlock()
	queue.length++
	log.Debugf("queue length: %v", queue.length)

}

// Pop 取出缓存的取证任务
func (queue *TaskQueue) Pop() *task {
	item := <-queue.taskCacheCh
	queueLock.Lock()
	defer queueLock.Unlock()
	queue.length--
	return item
}

// Result 取证结果
type Result struct {
	evidenceId  string //取证ID
	identityId  string //身份ID
	userId      string //用户ID
	userPrvKey  string //用户私钥
	filePath    string //文件绝对路径
	fileName    string //文件名称
	fileHash    []byte //文件哈希
	TxId        string //交易ID
	TxTime      string //交易时间
	para        *param.CreateObtainEvidenceParam
	taskType    string
	client      *client.IchainClient
	letterParam *param.EvidenceLetterParam //存证函参数
	err         error
}

// TaskManager 取证任务管理器
type TaskManager struct {
	items      map[string]*task //正在进行取证的任务，key为取证ID
	dbEngine   *xorm.Engine
	workerPool *ants.Pool //取证任务协程池
}

// InitTaskManager 初始化管理器
func InitTaskManager() (*TaskManager, error) {
	if taskMana != nil {
		return taskMana, nil
	}

	dbEngine := db.GetEngine()
	if dbEngine == nil {
		return nil, errors.New("[InitTaskManager]: failed to get dbEngine!")
	}

	err := dbEngine.Ping()
	if err != nil {
		log.Errorf("Failed to connect to the database: %v", err)
		return nil, err
	}
	//获取容量大小
	workerPool, err := ants.NewPool(config.GetInt("icep.cdp.task-capacity", 50), ants.WithPreAlloc(false))
	if err != nil {
		return nil, err
	}
	taskMana = &TaskManager{
		items:      make(map[string]*task),
		dbEngine:   dbEngine,
		workerPool: workerPool,
	}
	return taskMana, nil
}

//启动任务
func (mana *TaskManager) startTask() {
	//新建协程，添加任务
	resCh := make(chan *Result, config.GetInt("icep.cdp.task-capacity", 50))
	go func() {
		for {
			//取出待执行的任务
			wTask := GetTaskQueue().Pop()
			log.Debugf("pop task from the queue, evidenceId: %v, url: %v", wTask.evidenceId, wTask.para.ObtainUrl)
			//协程池空闲则执行
			err := mana.workerPool.Submit(
				func() {
					//添加任务
					taskMana.addTask(resCh, wTask)
				})
			if err != nil {
				log.Fatal(err)
			}
		}
	}()

	//新建协程监听执行结果
	go func() {
		for {
			//获取结果，执行后续逻辑
			select {
			case res := <-resCh:
				//某个任务完成后执行后续逻辑
				log.Debug("processing the task result")
				taskMana.handleTaskResult(res)
				log.Debug("the task result has been processed")
			}
		}
	}()
}

//新建协程，执行具体的取证逻辑
func (mana *TaskManager) addTask(ch chan<- *Result, tk *task) {
	//捕获 page load error
	defer func() {
		if err := recover(); err != nil {
			log.Error(err)
		}
	}()

	//记录任务
	manaLock.Lock()
	mana.items[tk.evidenceId] = tk
	manaLock.Unlock()

	//获取chrome context
	chromeContext, chromeContextCancel, err := GetChromeCtx()
	if err != nil {
		ch <- &Result{
			evidenceId: tk.evidenceId,
			identityId: tk.identityId,
			err:        err,
		}
		return
	}
	log.Debugf("get chrome context:%s", tk.evidenceId)

	//为context设置超时时间
	taskCtx, taskCancel := context.WithTimeout(context.Background(), time.Second*time.Duration(config.GetInt("icep.cdp.task-timeout", 180)))
	//执行任务并计时，超时强制终止
	taskTiming(taskCtx, chromeContext, ch, tk)

	taskCancel()
	log.Debugf("task context cancel:%s", tk.evidenceId)

	//关闭标签页
	chromeContextCancel()
	log.Debugf("chrome context cancel:%s", tk.evidenceId)
}

//执行任务并计时
func taskTiming(taskCtx context.Context, chromeCtx context.Context, resCh chan<- *Result, tk *task) {
	timingCh := make(chan bool, 1)
	go execTask(timingCh, chromeCtx, resCh, tk)

	select {
	//正常执行结束
	case <-timingCh:
		log.Infof("task has ended, evidenceId: %v", tk.evidenceId)

	//执行超时
	case <-taskCtx.Done():
		log.Infof("task execution timeout, evidenceId: %v", tk.evidenceId)
	}
}

//执行任务
func execTask(ch chan<- bool, ctx context.Context, resCh chan<- *Result, tk *task) {
	log.Debugf("exec task, ch: %v, ctx: %v, resCh: %v, tk: %v", ch != nil, ctx != nil, resCh != nil, tk != nil)

	if err := chromedp.Run(ctx, captureWebPage(resCh, tk)); err != nil {
		log.Infof("exec task err: %v", err)
		resCh <- &Result{
			evidenceId: tk.evidenceId,
			identityId: tk.identityId,
			userId:     tk.userId,
			userPrvKey: tk.userPrvKey,
			para:       tk.para,
			taskType:   tk.taskType,
			err:        err,
		}
	}
	ch <- true
}

//取证完成后删除任务
func (mana *TaskManager) done(evidenceId string) {
	manaLock.Lock()
	defer manaLock.Unlock()
	_, ok := mana.items[evidenceId]
	if !ok {
		log.Debugf("item not found, key: %v", evidenceId)
		return
	}
	log.Debugf("delete item form the task manager, evidenceId: %v", evidenceId)
	delete(mana.items, evidenceId)
}

//处理取证结果
func (mana *TaskManager) handleTaskResult(res *Result) {
	defer func() {
		if err := recover(); err != nil {
			log.Error("recover err:", err)
		}
	}()

	//判断是否成功抓取网页(将截图或文件存储至临时目录)
	if res.err != nil {
		log.Errorf("receive err: %v", res.err)
		mana.dealWithErr(res, res.err)
		return
	}

	//签名
	signature := gm.Sign(res.userPrvKey, res.fileHash)
	//保存签名到签名表
	err := InsertSignature(mana.dbEngine, res.evidenceId, signature)
	if err != nil {
		//向通道中写入结果
		mana.dealWithErr(res, err)
		return
	}
	log.Debugf("insert signature, evidence id: %v, signature: %v", res.evidenceId, signature)

	//上传云存储获得URL
	uploadClient := upload.GetUploadClient()
	//云存储上传路径为：取证ID/文件名
	url, err := uploadClient.UploadFile(res.filePath, res.evidenceId+"/"+res.fileName)

	if err != nil {
		mana.dealWithErr(res, err)
		return
	}
	log.Debugf("upload file, webPageURL:%v, fileName: %v, fileURL: %v", res.para.ObtainUrl, res.fileName, url)
	//上传云存储后删除文件
	defer func(path string) {
		err := os.RemoveAll(path)
		if err != nil {
			log.Errorf("Failed to remove temporary directory: %v", err)
		}
	}(res.filePath)

	//存入取证明细表
	err = InsertObtainEvidenceDetail(mana.dbEngine, res.evidenceId, res.fileName, string(res.fileHash), url)
	if err != nil {
		mana.dealWithErr(res, err)
		return
	}

	//上链
	res.TxId, res.TxTime, err = InsertChain(res.client, res.evidenceId, signature, string(res.fileHash), *res.para, res.fileName,
		string(res.fileHash), url)
	if err != nil {
		mana.dealWithErr(res, err)
		return
	}

	//生成存证函
	letterUrl, err := GeneAttestLetter(res.client, res.letterParam, res.evidenceId, res.TxId, res.TxTime,
		string(res.fileHash), res.taskType)
	if err != nil {
		mana.dealWithErr(res, err)
		return
	}

	//更新取证表
	err = UpdateObtainEvidenceData(mana.dbEngine, res.evidenceId, res.identityId, constant.ObtainStatusSuccess, letterUrl, res)
	if err != nil {
		mana.dealWithErr(res, err)
		return
	}
	log.Debugf("result has been returned, evidenceId: %v", res.evidenceId)
	//从任务管理器中清除任务
	mana.done(res.evidenceId)
}

//处理错误
func (mana *TaskManager) dealWithErr(res *Result, err error) {
	mana.done(res.evidenceId)
	err2 := UpdateObtainEvidenceData(mana.dbEngine, res.evidenceId, res.identityId, constant.ObtainStatusFailure, "", res)
	if err2 != nil {
		log.Errorf("update obtain evidence err: %s", err2.Error())
		return
	}
	log.Debug("update obtain evidence data success")
}

// GetChromeCtx 获取chrome context
func GetChromeCtx() (ctx context.Context, ctxCancel context.CancelFunc, err error) {
	//chrome已初始化，新建context
	if chromeCtx != nil {
		ctx, ctxCancel = chromedp.NewContext(chromeCtx)
		return
	}
	//chrome未初始化
	chromeCtx, chromeCtxCancel, err = initChrome()
	if err != nil {
		return
	}
	ctx, ctxCancel = chromedp.NewContext(chromeCtx)
	return
}

//初始化chrome
func initChrome() (context.Context, context.CancelFunc, error) {
	var options []chromedp.ExecAllocatorOption
	options = append(options, chromedp.Flag("headless", config.GetBool("icep.cdp.headless")))
	options = append(options, chromedp.Flag("allow-cross-origin-auth-prompt", ""))
	options = append(options, chromedp.Flag("ignore-certificate-errors", ""))
	options = append(options, chromedp.Flag("no-sandbox", ""))
	options = append(options, chromedp.Flag("disable-dev-shm-usage", ""))
	options = append(options, chromedp.Flag("window-size", config.Get("icep.cdp.window-size", "1920,1080")))
	options = append(options, chromedp.UserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/93.0.4577.82 Safari/537.36"))
	options = append(options, chromedp.DisableGPU)

	eCtx, eCancel := chromedp.NewExecAllocator(context.Background(), options...)
	nCtx, _ := chromedp.NewContext(eCtx)
	// 启动chrome
	if err := chromedp.Run(nCtx); err != nil {
		return nil, nil, errors.Errorf("init chrome err: %v", err)
	}
	return nCtx, eCancel, nil
}

//获取网页快照
func captureWebPage(ch chan<- *Result, tk *task) chromedp.Tasks {
	//验证URL
	url := VerifyUrl(tk)

	return chromedp.Tasks{
		//加载网页
		chromedp.Navigate(url),
		chromedp.ActionFunc(func(ctx context.Context) error {
			//待状态为complete时，执行后续逻辑
			err := getWebPageState(ctx)
			if err != nil {
				return err
			}

			//获取网页尺寸
			width, height, err := getWebPageSize(ctx)
			if err != nil {
				return err
			}
			//加载非固定高度的网页
			height, err = scrollToBottom(ctx, height)
			if err != nil {
				return err
			}

			//设定屏幕尺寸
			log.Debugf("set device metric, width: %v, height: %v", width, height)
			err = setDeviceMetrics(ctx, width, height)
			if err != nil {
				return err
			}

			//配置文件中读取网页加载延时时间
			loadDelay := config.GetInt("icep.cdp.load-delay", 2)
			//事件通道
			eventCh := make(chan bool)
			//监听事件
			go chromedp.ListenTarget(ctx, func(event interface{}) {
				//监听image资源请求事件
				switch eventType := event.(type) {
				case *network.EventRequestWillBeSent:
					if eventType.Type == "Image" {
						eventCh <- true
					}
				case *network.EventRequestWillBeSentExtraInfo:
					i, ok := eventType.Headers["Sec-Fetch-Dest"]
					if ok && i == "image" {
						eventCh <- true
					}
				}
			})
			//等待网页加载完成
			waitToLoad(loadDelay, eventCh)

			var filePath, fileName string
			var fileHash []byte
			switch tk.taskType {
			//网页文件
			case constant.ObtainTypeHtml:
				filePath, fileName, fileHash, err = captureSnapshot(ctx, tk)
				if err != nil {
					return err
				}
			//屏幕截图
			case constant.ObtainTypePict:
				filePath, fileName, fileHash, err = captureScreenShot(ctx, tk)
				if err != nil {
					return err
				}
			}
			//组合执行结果写入通道
			ch <- &Result{
				evidenceId:  tk.evidenceId,
				identityId:  tk.identityId,
				userId:      tk.userId,
				userPrvKey:  tk.userPrvKey,
				filePath:    filePath,
				fileName:    fileName,
				fileHash:    fileHash,
				para:        tk.para,
				taskType:    tk.taskType,
				client:      tk.client,
				letterParam: tk.letterParam,
				err:         nil,
			}

			//清除浏览器缓存
			network.ClearBrowserCache().Do(ctx)
			return nil
		}),
	}
}

//获取网页状态
func getWebPageState(ctx context.Context) error {
	for {
		result, _, err := evalAsValue(ctx, "document.readyState")
		if err != nil {
			return err
		}
		res := strings.Trim(string(result.Value), "\"")
		if res == "complete" {
			break
		}
		time.Sleep(time.Millisecond * 100)
	}
	return nil
}

//将页面滚动至底部
func scrollToBottom(ctx context.Context, height int) (int, error) {
	for {
		//将页面滚动至最底部
		_, _, err := evalAsValue(ctx, fmtScrollExpr(height))
		if err != nil {
			return height, err
		}
		time.Sleep(time.Millisecond * 500)
		//再次获取网页尺寸，判断与上次是否相等
		_, curHeight, err := getWebPageSize(ctx)
		if err != nil {
			return height, err
		}
		log.Debugf("current height of the webpage: %v", curHeight)
		//高度相等说明到达页面最底部
		if curHeight == height {
			log.Debug("reached the bottom of the webpage")
			break
		}
		height = curHeight
	}
	return height, nil
}

//获取网页尺寸
func getWebPageSize(ctx context.Context) (width, height int, err error) {
	_, _, _, _, _, cssContentSize, err := page.GetLayoutMetrics().Do(ctx)
	if err != nil {
		return
	}
	width = int(math.Ceil(cssContentSize.Width))
	height = int(math.Ceil(cssContentSize.Height))
	return
}

// 设置屏幕尺寸
func setDeviceMetrics(ctx context.Context, width, height int) error {
	err := emulation.SetDeviceMetricsOverride(int64(width), int64(height), 1, false).Do(ctx)
	if err != nil {
		return errors.Errorf("[setDeviceMetrics] err: %v", err)
	}
	return nil
}

//等待网页加载
func waitToLoad(delay int, eventCh chan bool) {
	//网页加载计时，超时后强制结束
	for {
		select {
		//网页加载延时
		case <-time.After(time.Second * time.Duration(delay)):
			log.Debug("webpage has been loaded.")
			return
		case <-eventCh:
			break
		}
	}
}

//执行js表达式
func evalAsValue(ctx context.Context, expression string) (*runtime.RemoteObject, *runtime.ExceptionDetails, error) {
	result, detail, err := chromedp.EvalAsValue(&runtime.EvaluateParams{
		Expression:    expression,
		ReturnByValue: true,
		Timeout:       2000,
	}).Do(ctx)
	if err != nil {
		return nil, nil, errors.Errorf("[evalAsValue] err: %v", err)
	}

	return result, detail, nil
}

//获取网页截图png
func captureScreenShot(ctx context.Context, tk *task) (filePath, fileName string, fileHash []byte, err error) {
	dataBytes, err := page.CaptureScreenshot().WithFormat(page.CaptureScreenshotFormatPng).Do(ctx)
	if err != nil {
		return
	}

	path, err := creatTempDir(tk.evidenceId)
	if err != nil {
		return
	}
	log.Debugf("creat temporary directory, userId: %v, evidenceId:%v, path: %v", tk.userId, tk.evidenceId, path)
	//定义文件名（取证ID.png)
	fileName = "webPage.png"
	log.Debugf("file name: %v", fileName)

	filePath = filepath.Join(path, fileName)
	err = ioutil.WriteFile(filePath, dataBytes, 0644)
	if err != nil {
		return
	}
	_, _, fileHash, err = GetDataHash(dataBytes)
	if err != nil {
		return
	}
	return
}

//获取网页文件mhtml
func captureSnapshot(ctx context.Context, tk *task) (filePath, fileName string, fileHash []byte, err error) {
	mhtml, err := page.CaptureSnapshot().WithFormat(page.CaptureSnapshotFormatMhtml).Do(ctx)
	if err != nil {
		return
	}
	path, err := creatTempDir(tk.evidenceId)
	if err != nil {
		return
	}
	log.Debugf("creat temporary directory, userId: %v, evidenceId:%v, path: %v", tk.userId, tk.evidenceId, path)
	//定义文件名（取证ID.mhtml)
	fileName = "webPage.mhtml"
	filePath = filepath.Join(path, fileName)
	err = ioutil.WriteFile(filePath, []byte(mhtml), 0644)
	if err != nil {
		return
	}
	_, _, fileHash, err = GetDataHash([]byte(mhtml))
	if err != nil {
		return
	}
	return
}

// GeneAttestLetter 生成存证函
func GeneAttestLetter(client *client.IchainClient, letterParam *param.EvidenceLetterParam, evidenceId, txId, txTime,
	fileHash, taskType string) (string, error) {

	blockInfo := client.GetBlockInfoById(txId)
	if blockInfo == nil {
		return "", errors.New("get Block Info err!")
	}

	if taskType == constant.ObtainTypeHtml {
		letterParam.EvidenceType = "网页取证"
	} else if taskType == constant.ObtainTypePict {
		letterParam.EvidenceType = "截屏取证"
	} else if taskType == constant.ObtainTypeVideo {
		letterParam.EvidenceType = "录屏取证"
	}

	letterParam.EvidenceId = evidenceId
	letterParam.TxId = txId
	letterParam.EvidenceTime = txTime
	letterParam.BlockNum = int(blockInfo.BlockNum)
	letterParam.BlockHash = blockInfo.BlockHash
	letterParam.FileHash = fileHash

	tempDir, err := creatTempDir(evidenceId)
	if err != nil {
		return "", err
	}
	//上传云存储后删除
	defer func(path string) {
		err = os.RemoveAll(path)
		if err != nil {
			log.Errorf("Failed to remove temporary directory: %v", err)
		}
	}(tempDir)

	path := filepath.Join(tempDir, "存证函.png")
	letterParam.TargetFile = path

	err = utils.CreateEvidenceLetter(letterParam)
	if err != nil {
		log.Errorf("create evidence letter error: %v", err)
		return "", err
	}

	fileUrl, err := upload.GetUploadClient().UploadFile(path, strings.Join([]string{evidenceId, "存证函.png"}, "/"))
	if err != nil {
		log.Errorf("upload file error: %v", err)
		return "", err
	}
	return fileUrl, nil
}
